use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
use std::ptr;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

/// 高性能无锁哈希表，专门为GC场景优化
/// 使用分段锁减少竞争，提高并发性能
pub struct OptimizedLockFreeHashMap<K, V>
where
    K: Eq + Clone + Hash,
    V: Clone,
{
    /// 分段存储，每个段有自己的锁
    segments: Vec<Segment<K, V>>,
    /// 段的数量，通常是2的幂
    segment_shift: usize,
    /// 段掩码，用于快速计算段索引
    segment_mask: usize,
    /// 元素总数
    len: AtomicUsize,
}

/// 哈希表的一个段
struct Segment<K, V>
where
    K: Eq + Clone + Hash,
    V: Clone,
{
    /// 桶数组
    buckets: Vec<AtomicPtr<Entry<K, V>>>,
    /// 段内元素数量
    len: AtomicUsize,
    /// 段锁，用于扩容等操作
    lock: AtomicPtr<()>,
}

/// 哈希表条目
struct Entry<K, V> {
    key: K,
    value: V,
    hash: u64,
    next: AtomicPtr<Entry<K, V>>,
}

impl<K, V> OptimizedLockFreeHashMap<K, V>
where
    K: Eq + Clone + Hash,
    V: Clone,
{
    /// 创建新的哈希表
    pub fn new() -> Self {
        Self::with_capacity(16, 16)
    }
    
    /// 创建指定容量和段数的哈希表
    pub fn with_capacity(initial_capacity: usize, segments: usize) -> Self {
        // 确保段数是2的幂
        let segment_shift = 64 - (segments as u64).leading_zeros() as usize;
        let segment_mask = segments - 1;
        
        let mut segs = Vec::with_capacity(segments);
        for _ in 0..segments {
            let bucket_count = (initial_capacity / segments).max(4);
            let mut buckets = Vec::with_capacity(bucket_count);
            for _ in 0..bucket_count {
                buckets.push(AtomicPtr::new(ptr::null_mut()));
            }
            
            segs.push(Segment {
                buckets,
                len: AtomicUsize::new(0),
                lock: AtomicPtr::new(ptr::null_mut()),
            });
        }
        
        OptimizedLockFreeHashMap {
            segments: segs,
            segment_shift,
            segment_mask,
            len: AtomicUsize::new(0),
        }
    }
    
    /// 计算键的哈希值
    fn hash_key(&self, key: &K) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }
    
    /// 根据哈希值计算段索引
    fn segment_for_hash(&self, hash: u64) -> usize {
        (hash >> self.segment_shift) as usize & self.segment_mask
    }
    
    /// 根据哈希值计算桶索引
    fn bucket_for_hash(&self, hash: u64, bucket_count: usize) -> usize {
        (hash as usize) & (bucket_count - 1)
    }
    
    /// 插入键值对
    pub fn insert(&self, key: K, value: V) {
        let hash = self.hash_key(&key);
        let segment_idx = self.segment_for_hash(hash);
        let segment = &self.segments[segment_idx];
        
        // 计算桶索引
        let bucket_count = segment.buckets.len();
        let bucket_idx = self.bucket_for_hash(hash, bucket_count);
        
        // 创建新条目
        let new_entry = Box::into_raw(Box::new(Entry {
            key: key.clone(),
            value,
            hash,
            next: AtomicPtr::new(ptr::null_mut()),
        }));
        
        // 尝试插入到桶中
        loop {
            let mut current = segment.buckets[bucket_idx].load(Ordering::Acquire);
            let mut prev: *mut Entry<K, V> = ptr::null_mut();
            
            // 检查是否已存在相同的键
            while !current.is_null() {
                let entry = unsafe { &*current };
                if entry.hash == hash && entry.key == key {
                    // 键已存在，更新值
                    unsafe {
                        ptr::write(&mut (*current).value, (*new_entry).value.clone());
                        let _ = Box::from_raw(new_entry);
                    }
                    return;
                }
                prev = current;
                current = entry.next.load(Ordering::Acquire);
            }
            
            // 插入新条目
            unsafe { (*new_entry).next.store(current, Ordering::Relaxed); }
            
            let success = if prev.is_null() {
                // 插入到桶头
                segment.buckets[bucket_idx].compare_exchange(
                    current, new_entry, 
                    Ordering::Release, Ordering::Relaxed
                ).is_ok()
            } else {
                // 插入到链表中
                unsafe {
                    (*prev).next.compare_exchange(
                        current, new_entry,
                        Ordering::Release, Ordering::Relaxed
                    ).is_ok()
                }
            };
            
            if success {
                segment.len.fetch_add(1, Ordering::Relaxed);
                self.len.fetch_add(1, Ordering::Relaxed);
                break;
            }
            
            // CAS失败，重试
        }
    }
    
    /// 获取值
    pub fn get(&self, key: &K) -> Option<V> {
        let hash = self.hash_key(key);
        let segment_idx = self.segment_for_hash(hash);
        let segment = &self.segments[segment_idx];
        
        let bucket_count = segment.buckets.len();
        let bucket_idx = self.bucket_for_hash(hash, bucket_count);
        
        let mut current = segment.buckets[bucket_idx].load(Ordering::Acquire);
        
        while !current.is_null() {
            let entry = unsafe { &*current };
            if entry.hash == hash && entry.key == *key {
                return Some(entry.value.clone());
            }
            current = entry.next.load(Ordering::Acquire);
        }
        
        None
    }
    
    /// 删除键值对
    pub fn delete(&self, key: &K) -> Option<V> {
        let hash = self.hash_key(key);
        let segment_idx = self.segment_for_hash(hash);
        let segment = &self.segments[segment_idx];
        
        let bucket_count = segment.buckets.len();
        let bucket_idx = self.bucket_for_hash(hash, bucket_count);
        
        let mut current = segment.buckets[bucket_idx].load(Ordering::Acquire);
        let mut prev: *mut Entry<K, V> = ptr::null_mut();
        
        while !current.is_null() {
            let entry = unsafe { &*current };
            if entry.hash == hash && entry.key == *key {
                let next = entry.next.load(Ordering::Acquire);
                
                let success = if prev.is_null() {
                    // 删除桶头
                    segment.buckets[bucket_idx].compare_exchange(
                        current, next,
                        Ordering::Release, Ordering::Relaxed
                    ).is_ok()
                } else {
                    // 删除链表中的节点
                    unsafe {
                        (*prev).next.compare_exchange(
                            current, next,
                            Ordering::Release, Ordering::Relaxed
                        ).is_ok()
                    }
                };
                
                if success {
                    segment.len.fetch_sub(1, Ordering::Relaxed);
                    self.len.fetch_sub(1, Ordering::Relaxed);
                    return Some(unsafe { Box::from_raw(current).value });
                }
                
                // CAS失败，重新开始
                current = segment.buckets[bucket_idx].load(Ordering::Acquire);
                prev = ptr::null_mut();
                continue;
            }
            
            prev = current;
            current = entry.next.load(Ordering::Acquire);
        }
        
        None
    }
    
    /// 获取元素数量
    pub fn len(&self) -> usize {
        self.len.load(Ordering::Relaxed)
    }
    
    /// 检查是否为空
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    
    /// 批量操作，减少锁竞争
    pub fn batch_insert(&self, pairs: &[(K, V)]) {
        // 按段分组，减少跨段操作
        let mut segment_groups = std::collections::HashMap::new();
        
        for (key, value) in pairs {
            let hash = self.hash_key(key);
            let segment_idx = self.segment_for_hash(hash);
            segment_groups.entry(segment_idx).or_insert_with(Vec::new).push((key.clone(), value.clone(), hash));
        }
        
        // 按段处理
        for (segment_idx, items) in segment_groups {
            let segment = &self.segments[segment_idx];
            for (key, value, hash) in items {
                // 计算桶索引
                let bucket_count = segment.buckets.len();
                let bucket_idx = self.bucket_for_hash(hash, bucket_count);
                
                // 创建新条目
                let new_entry = Box::into_raw(Box::new(Entry {
                    key: key.clone(),
                    value,
                    hash,
                    next: AtomicPtr::new(ptr::null_mut()),
                }));
                
                // 尝试插入到桶中
                loop {
                    let mut current = segment.buckets[bucket_idx].load(Ordering::Acquire);
                    let mut prev: *mut Entry<K, V> = ptr::null_mut();
                    
                    // 检查是否已存在相同的键
                    while !current.is_null() {
                        let entry = unsafe { &*current };
                        if entry.hash == hash && entry.key == key {
                            // 键已存在，更新值
                            unsafe {
                                ptr::write(&mut (*current).value, (*new_entry).value.clone());
                                let _ = Box::from_raw(new_entry);
                            }
                            break;
                        }
                        prev = current;
                        current = entry.next.load(Ordering::Acquire);
                    }
                    
                    // 如果键不存在，插入新条目
                    if current.is_null() {
                        unsafe { (*new_entry).next.store(ptr::null_mut(), Ordering::Relaxed); }
                        
                        let success = if prev.is_null() {
                            // 插入到桶头
                            segment.buckets[bucket_idx].compare_exchange(
                                current, new_entry, 
                                Ordering::Release, Ordering::Relaxed
                            ).is_ok()
                        } else {
                            // 插入到链表中
                            unsafe {
                                (*prev).next.compare_exchange(
                                    current, new_entry,
                                    Ordering::Release, Ordering::Relaxed
                                ).is_ok()
                            }
                        };
                        
                        if success {
                            segment.len.fetch_add(1, Ordering::Relaxed);
                            self.len.fetch_add(1, Ordering::Relaxed);
                            break;
                        }
                        
                        // CAS失败，重试
                    } else {
                        // 键已存在，已更新，跳出循环
                        break;
                    }
                }
            }
        }
    }
    
    /// 迭代器，用于遍历所有键值对
    pub fn iter(&self) -> OptimizedHashMapIter<'_, K, V> {
        OptimizedHashMapIter {
            map: self,
            segment_idx: 0,
            bucket_idx: 0,
            current: ptr::null_mut(),
        }
    }
}

/// 迭代器实现
pub struct OptimizedHashMapIter<'a, K, V>
where
    K: Eq + Clone + Hash,
    V: Clone,
{
    map: &'a OptimizedLockFreeHashMap<K, V>,
    segment_idx: usize,
    bucket_idx: usize,
    current: *const Entry<K, V>,
}

impl<'a, K, V> Iterator for OptimizedHashMapIter<'a, K, V>
where
    K: Eq + Clone + Hash,
    V: Clone,
{
    type Item = (K, V);
    
    fn next(&mut self) -> Option<Self::Item> {
        // 如果当前节点不为空，尝试获取下一个节点
        if !self.current.is_null() {
            let entry = unsafe { &*self.current };
            let next = entry.next.load(Ordering::Acquire);
            if !next.is_null() {
                self.current = next;
                return Some((entry.key.clone(), entry.value.clone()));
            }
        }
        
        // 遍历所有段和桶
        while self.segment_idx < self.map.segments.len() {
            let segment = &self.map.segments[self.segment_idx];
            
            // 遍历当前段的所有桶
            while self.bucket_idx < segment.buckets.len() {
                // 如果当前节点为空，获取桶的第一个节点
                if self.current.is_null() {
                    self.current = segment.buckets[self.bucket_idx].load(Ordering::Acquire);
                }
                
                if !self.current.is_null() {
                    let entry = unsafe { &*self.current };
                    let result = (entry.key.clone(), entry.value.clone());
                    
                    // 移动到下一个节点
                    let next = entry.next.load(Ordering::Acquire);
                    self.current = next;
                    
                    // 如果没有下一个节点，移动到下一个桶
                    if next.is_null() {
                        self.bucket_idx += 1;
                        self.current = ptr::null_mut();
                    }
                    
                    return Some(result);
                }
                
                self.bucket_idx += 1;
            }
            
            // 移动到下一个段
            self.segment_idx += 1;
            self.bucket_idx = 0;
            self.current = ptr::null_mut();
        }
        
        None
    }
}

impl<K, V> Drop for OptimizedLockFreeHashMap<K, V>
where
    K: Eq + Clone + Hash,
    V: Clone,
{
    fn drop(&mut self) {
        // 清理所有条目
        for segment in &self.segments {
            for bucket in &segment.buckets {
                let mut current = bucket.load(Ordering::Relaxed);
                while !current.is_null() {
                    let entry = unsafe { Box::from_raw(current) };
                    current = entry.next.load(Ordering::Relaxed);
                }
            }
        }
    }
}